﻿#define WIN32_LEAN_AND_MEAN  // 排除极少使用的内容
#include <winsock2.h>
#include <windows.h>
#include <ws2tcpip.h>
#include <iostream>
#include <vector>
#include <thread>
#include <atomic>
#include <cstring>

#pragma comment(lib, "ws2_32.lib")

// 禁用 Winsock 1 相关的警告
#define _WINSOCK_DEPRECATED_NO_WARNINGS
// 防止 Windows.h 包含 min/max 宏
#define NOMINMAX
const int BUFFER_SIZE = 4096;     // 网络缓冲区大小
const char* DEFAULT_PORT = "8080"; // 服务器端口

// 消息类型
enum MessageType {
    SREEN_INFO = 1,
    SREEN_DATA,
    MOU_EVENT,
    BOARD_EVENT
};

// 屏幕信息结构
struct ScreenInfo {
    int width;
    int height;
    int bpp;  // 位深度
};

// 鼠标事件结构
struct MouseEvent {
    int x;
    int y;
    int eventType;  // 0:移动, 1:左键按下, 2:左键释放
};

// 键盘事件结构
struct KeyboardEvent {
    unsigned int keyCode;
    bool isDown;
};

std::atomic<bool> running(true);
SOCKET serverSocket = INVALID_SOCKET;
HWND hwnd = NULL;
HDC hdc = NULL;
HBITMAP hBitmap = NULL;
HDC hdcMem = NULL;
BITMAPINFOHEADER bi = { 0 };
ScreenInfo screenInfo;
bool screenInfoReceived = false;

// 窗口过程函数
LRESULT CALLBACK WindowProc(HWND hwnd, UINT uMsg, WPARAM wParam, LPARAM lParam) {
    switch (uMsg) {
    case WM_DESTROY:
        PostQuitMessage(0);
        running = false;
        return 0;

    case WM_PAINT: {
        PAINTSTRUCT ps;
        HDC hdc = BeginPaint(hwnd, &ps);

        if (screenInfoReceived && hBitmap) {
            RECT rect;
            GetClientRect(hwnd, &rect);

            // 使用内存DC中的位图进行绘制
            BitBlt(hdc, 0, 0, rect.right, rect.bottom, hdcMem, 0, 0, SRCCOPY);
        }

        EndPaint(hwnd, &ps);
        return 0;
    }

    case WM_LBUTTONDOWN: {
        int x = LOWORD(lParam);
        int y = HIWORD(lParam);

        RECT rect;
        GetClientRect(hwnd, &rect);

        // 缩放坐标到服务器屏幕尺寸
        int serverX = (x * screenInfo.width) / rect.right;
        int serverY = (y * screenInfo.height) / rect.bottom;

        // 发送鼠标按下事件
        MouseEvent mouseEvent = { serverX, serverY, 1 };
        int messageType = MOUSE_EVENT;
        send(serverSocket, (char*)&messageType, sizeof(messageType), 0);
        send(serverSocket, (char*)&mouseEvent, sizeof(mouseEvent), 0);
        return 0;
    }

    case WM_LBUTTONUP: {
        int x = LOWORD(lParam);
        int y = HIWORD(lParam);

        RECT rect;
        GetClientRect(hwnd, &rect);

        // 缩放坐标到服务器屏幕尺寸
        int serverX = (x * screenInfo.width) / rect.right;
        int serverY = (y * screenInfo.height) / rect.bottom;

        // 发送鼠标释放事件
        MouseEvent mouseEvent = { serverX, serverY, 2 };
        int messageType = MOUSE_EVENT;
        send(serverSocket, (char*)&messageType, sizeof(messageType), 0);
        send(serverSocket, (char*)&mouseEvent, sizeof(mouseEvent), 0);
        return 0;
    }

    case WM_MOUSEMOVE: {
        if (wParam & MK_LBUTTON) {
            int x = LOWORD(lParam);
            int y = HIWORD(lParam);

            RECT rect;
            GetClientRect(hwnd, &rect);

            // 缩放坐标到服务器屏幕尺寸
            int serverX = (x * screenInfo.width) / rect.right;
            int serverY = (y * screenInfo.height) / rect.bottom;

            // 发送鼠标移动事件
            MouseEvent mouseEvent = { serverX, serverY, 0 };
            int messageType = MOUSE_EVENT;
            send(serverSocket, (char*)&messageType, sizeof(messageType), 0);
            send(serverSocket, (char*)&mouseEvent, sizeof(mouseEvent), 0);
        }
        return 0;
    }

    case WM_KEYDOWN: {
        // 发送键盘按下事件
        KeyboardEvent keyboardEvent = { (unsigned int)wParam, true };
        int messageType = BOARD_EVENT;
        send(serverSocket, (char*)&messageType, sizeof(messageType), 0);
        send(serverSocket, (char*)&keyboardEvent, sizeof(keyboardEvent), 0);
        return 0;
    }
    case WM_KEYUP: {
        // 发送键盘释放事件
        KeyboardEvent keyboardEventUp = { (unsigned int)wParam, false };
        int messageTypeUp = BOARD_EVENT;
        send(serverSocket, (char*)&messageTypeUp, sizeof(messageTypeUp), 0);
        send(serverSocket, (char*)&keyboardEventUp, sizeof(keyboardEventUp), 0);
        return 0;
    }
    }

    return DefWindowProc(hwnd, uMsg, wParam, lParam);
}

// 接收和处理屏幕数据
// 接收和处理屏幕数据
void ReceiveScreenData() {
    char buffer[BUFFER_SIZE];
    int bytesReceived;

    while (running) {
        // 接收消息类型
        if ((bytesReceived = recv(serverSocket, buffer, sizeof(int), 0)) <= 0) {
            std::cerr << "Connection closed by server." << std::endl;
            running = false;
            break;
        }

        int messageType = *(int*)buffer;

        switch (messageType) {
        case SREEN_INFO: {
            if ((bytesReceived = recv(serverSocket, buffer, sizeof(ScreenInfo), 0)) <= 0) {
                running = false;
                break;
            }

            screenInfo = *(ScreenInfo*)buffer;
            std::cout << "Screen info received: " << screenInfo.width << "x" << screenInfo.height
                << ", " << screenInfo.bpp << "bpp" << std::endl;

            // 初始化位图信息
            bi.biSize = sizeof(BITMAPINFOHEADER);
            bi.biWidth = screenInfo.width;
            // 将 biHeight 设置为正数
            bi.biHeight = screenInfo.height;
            bi.biPlanes = 1;
            bi.biBitCount = screenInfo.bpp;
            bi.biCompression = BI_RGB;

            // 删除旧位图
            if (hBitmap) {
                DeleteObject(hBitmap);
                hBitmap = NULL;
            }

            screenInfoReceived = true;

            // 调整窗口大小
            RECT rect = { 0, 0, screenInfo.width, screenInfo.height };
            AdjustWindowRect(&rect, WS_OVERLAPPEDWINDOW, FALSE);
            SetWindowPos(hwnd, NULL, 0, 0, rect.right - rect.left, rect.bottom - rect.top,
                SWP_NOMOVE | SWP_NOZORDER);

            InvalidateRect(hwnd, NULL, FALSE);
            break;
        }
        case SREEN_DATA: {
            if ((bytesReceived = recv(serverSocket, buffer, sizeof(int), 0)) <= 0) {
                running = false;
                break;
            }

            int dataSize = *(int*)buffer;
            std::vector<char> screenData(dataSize);

            int totalBytesReceived = 0;
            while (totalBytesReceived < dataSize) {
                bytesReceived = recv(serverSocket, screenData.data() + totalBytesReceived,
                    dataSize - totalBytesReceived, 0);
                if (bytesReceived <= 0) {
                    running = false;
                    break;
                }
                totalBytesReceived += bytesReceived;
            }

            if (!running) break;

            // 创建或更新位图
            if (!hBitmap) {
                hBitmap = CreateCompatibleBitmap(hdc, screenInfo.width, screenInfo.height);
                hdcMem = CreateCompatibleDC(hdc);
                SelectObject(hdcMem, hBitmap);
            }

            // 更新位图数据
            SetDIBits(hdcMem, hBitmap, 0, screenInfo.height, screenData.data(),
                (BITMAPINFO*)&bi, DIB_RGB_COLORS);

            // 请求重绘
            InvalidateRect(hwnd, NULL, FALSE);
            break;
        }
        }
    }
}

int main() {
    WSADATA wsaData;
    int iResult;

    // 初始化 Winsock
    iResult = WSAStartup(MAKEWORD(2, 2), &wsaData);
    if (iResult != 0) {
        std::cerr << "WSAStartup failed: " << iResult << std::endl;
        return 1;
    }

    // 创建窗口类
    WNDCLASSEX wc = { 0 };
    wc.cbSize = sizeof(WNDCLASSEX);
    wc.style = CS_HREDRAW | CS_VREDRAW;
    wc.lpfnWndProc = WindowProc;
    wc.hInstance = GetModuleHandle(NULL);
    wc.hCursor = LoadCursor(NULL, IDC_ARROW);
    wc.hbrBackground = (HBRUSH)(COLOR_WINDOW + 1);
    wc.lpszClassName = L"RemoteDesktopClient";

    if (!RegisterClassEx(&wc)) {
        std::cerr << "Window registration failed." << std::endl;
        WSACleanup();
        return 1;
    }

    // 创建窗口
    hwnd = CreateWindowEx(
        0,
        L"RemoteDesktopClient",
        L"Remote Desktop Client",
        WS_OVERLAPPEDWINDOW,
        CW_USEDEFAULT, CW_USEDEFAULT, 800, 600,
        NULL, NULL, GetModuleHandle(NULL), NULL
    );

    if (!hwnd) {
        std::cerr << "Window creation failed." << std::endl;
        WSACleanup();
        return 1;
    }

    hdc = GetDC(hwnd);

    // 显示窗口
    ShowWindow(hwnd, SW_SHOW);
    UpdateWindow(hwnd);

    // 连接到服务器
    struct addrinfo* result = NULL, * ptr = NULL, hints;

    ZeroMemory(&hints, sizeof(hints));
    hints.ai_family = AF_INET;
    hints.ai_socktype = SOCK_STREAM;
    hints.ai_protocol = IPPROTO_TCP;

    // 获取服务器地址
    char serverIP[100];
    std::cout << "Enter server IP: ";
    std::cin >> serverIP;

    // 解析服务器地址和端口
    iResult = getaddrinfo(serverIP, DEFAULT_PORT, &hints, &result);
    if (iResult != 0) {
        std::cerr << "getaddrinfo failed: " << iResult << std::endl;
        DestroyWindow(hwnd);
        WSACleanup();
        return 1;
    }

    // 尝试连接
    for (ptr = result; ptr != NULL; ptr = ptr->ai_next) {
        serverSocket = socket(ptr->ai_family, ptr->ai_socktype, ptr->ai_protocol);
        if (serverSocket == INVALID_SOCKET) {
            std::cerr << "socket failed: " << WSAGetLastError() << std::endl;
            continue;
        }

        iResult = connect(serverSocket, ptr->ai_addr, (int)ptr->ai_addrlen);
        if (iResult == SOCKET_ERROR) {
            closesocket(serverSocket);
            serverSocket = INVALID_SOCKET;
            continue;
        }
        break;
    }

    freeaddrinfo(result);

    if (serverSocket == INVALID_SOCKET) {
        std::cerr << "Unable to connect to server!" << std::endl;
        DestroyWindow(hwnd);
        WSACleanup();
        return 1;
    }

    std::cout << "Connected to server." << std::endl;

    // 创建线程接收屏幕数据
    std::thread screenThread(ReceiveScreenData);

    // 消息循环
    MSG msg = { 0 };
    while (running && GetMessage(&msg, NULL, 0, 0)) {
        TranslateMessage(&msg);
        DispatchMessage(&msg);
    }

    // 等待线程结束
    if (screenThread.joinable()) screenThread.join();

    // 清理资源
    if (hBitmap) DeleteObject(hBitmap);
    if (hdcMem) DeleteDC(hdcMem);
    if (hdc) ReleaseDC(hwnd, hdc);
    if (hwnd) DestroyWindow(hwnd);
    if (serverSocket != INVALID_SOCKET) closesocket(serverSocket);
    WSACleanup();

    return 0;
}